"""
通用文字识别
"""
import re
import fitz
import requests
import base64
from PIL import Image
from io import BytesIO
import argparse


def get_access_token():
    """请求token"""
    appid = 'c0bj142mbMiSR1tcavFG5UXu'
    app_secret = 'SV65kCcWzj5KYIcWkkBxyMyKwpGVfFYy'
    host = f'https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id={appid}&client_secret={app_secret}'
    response = requests.get(host)
    access_token = response.json()['access_token']
    return access_token
    

def get_trans_data(img_base64):
    """获取数据"""
    request_url_inv = "https://aip.baidubce.com/rest/2.0/ocr/v1/vat_invoice"  # 发票识别
    access_token = get_access_token()
    request_url = request_url_inv + '?access_token=' + access_token
    headers = {'content-type': 'application/x-www-form-urlencoded'}
    response = requests.post(request_url, data={'image': img_base64}, headers=headers)
    return response.json()['words_result']


def get_inv_data(path, debug=False):
    """获取发票信息"""
    if path.endswith(('.pdf', '.PDF')):
        doc = fitz.Document(path)
        page = doc[0]
        pix = page.get_pixmap()
        img = Image.frombytes('RGB', [pix.width, pix.height], pix.samples)
    else:
        img = Image.open(path)
    width, height = img.size
    min_width = 1000
    min_height = int(min_width / width * height)
    img = img.convert('L').resize((min_width, min_height))
    output_buffer = BytesIO()
    img.save(output_buffer, format='JPEG')
    byte_data = output_buffer.getvalue()
    base64_str = base64.b64encode(byte_data)
    data = get_trans_data(base64_str)
    if debug:
        print('==========百度接口数据：')
        for k, v in data.items():
            print(k, v)
    inv_code = data['InvoiceCode']
    inv_num = data['InvoiceNum']
    inv_chk = data['CheckCode']
    inv_money = data['TotalAmount']
    inv_date = re.sub(r'[年月日]', '', data['InvoiceDate'])
    return inv_code, inv_num, inv_date, inv_chk, inv_money
    

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-p', '--path', type=str, required=True, help='输入文件路径')
    parser.add_argument('-d', '--debug', action='store_true', help='是否打印OCR接口返回数据')
    args = parser.parse_args()
    data = get_inv_data(args.path, args.debug)
    print(data)
